import av
import numpy as np
import cv2

container = av.open("01.mp4", mode="w")
stream = container.add_stream("h264", rate=25)
stream.width = 640
stream.height = 480
stream.pix_fmt = "yuv420p"

for i in range(250):
    # 1. 生成彩色渐变背景，横向渐变
    img = np.zeros((stream.height, stream.width, 3), dtype=np.uint8)
    for x in range(stream.width):
        color = (
            x * 255 // stream.width,
            255 - x * 255 // stream.width,
            (x * 5 + i * 2) % 256,
        )
        img[:, x, :] = color

    # 2. 动态红色矩形
    top_left = (i * 6 % stream.width, 100)
    bottom_right = (top_left[0] + 100, 200)
    cv2.rectangle(img, top_left, bottom_right, (0, 0, 255), -1)  # BGR 红色填充

    # 3. 动态绿色圆圈
    center = (stream.width - (i * 8 % stream.width), 350)
    cv2.circle(img, center, 50, (0, 255, 0), -1)  # BGR 绿色填充

    # 4. 左上角写帧编号，白色文字
    cv2.putText(
        img, f"Frame {i}", (10, 40), cv2.FONT_HERSHEY_SIMPLEX, 1.2, (255, 255, 255), 2
    )

    frame = av.VideoFrame.from_ndarray(img, format="bgr24")
    packet = stream.encode(frame)
    if packet:
        container.mux(packet)

# Flush
packet = stream.encode(None)
if packet:
    container.mux(packet)

container.close()
